package com.ruoyi.common.utils.video;

import lombok.extern.slf4j.Slf4j;
import org.bytedeco.javacv.FFmpegFrameGrabber;
import org.bytedeco.javacv.Frame;
import org.bytedeco.javacv.Java2DFrameConverter;

import javax.imageio.ImageIO;
import java.awt.*;
import java.awt.geom.AffineTransform;
import java.awt.image.BufferedImage;
import java.io.File;

@Slf4j
public class VideoFrameExtractor {

    public static void main(String[] args) {
        String videoPath = "F:\\抽帧\\VID_20250621_133014.mp4"; // 输入视频路径
        String outputDir = "F:\\抽帧\\抽帧图片"; // 输出目录
        int frameInterval = 300; // 每隔多少帧抽取一帧（取决于视频帧率）

        try {
            extractFrames(videoPath, outputDir, frameInterval);
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    public static int extractFrames(String videoPath, String outputDir, int frameInterval) throws Exception {
        int originalNum = 0;
        FFmpegFrameGrabber grabber = new FFmpegFrameGrabber(videoPath);
        grabber.start();

        // 确保输出目录存在
        File outputFolder = new File(outputDir);
        if (!outputFolder.exists()) {
            outputFolder.mkdirs();
        }

        Java2DFrameConverter converter = new Java2DFrameConverter();
        int totalFrames = grabber.getLengthInFrames();
        log.info("视频总帧数: " + totalFrames);

        double displayRotation = grabber.getDisplayRotation();
        log.info("视频旋转角度: " + displayRotation + "度");

        for (int i = 0; i < totalFrames; i++) {
            Frame frame = grabber.grabImage();
            if (frame == null) {
                continue;
            }

            // 按间隔抽取帧
            if (i % frameInterval == 0) {
                BufferedImage image = converter.getBufferedImage(frame);
                if (image != null) {
                    // 根据旋转角度矫正图像
                    //BufferedImage rotatedImage = rotateImage(image, displayRotation);

                    String outputPath = String.format("%s/frame_%05d.png", outputDir, i);
                    ImageIO.write(image, "png", new File(outputPath));
                    // log.info("已保存: " + outputPath);
                    originalNum += 1;
                }
            }
        }

        grabber.stop();
        log.info("视频抽帧完成！");
        return originalNum;
    }

    /**
     * 根据指定角度旋转图像（最终解决方案）
     *
     * @param image 原始图像
     * @param angle 旋转角度（度），FFmpeg返回的角度
     * @return 旋转后的图像
     */
    private static BufferedImage rotateImage(BufferedImage image, double angle) {
        if(angle == 0) return image;

        // 将负角度转换为等效正角度（例如 -90° → 270°）
        if (angle < 0) {
            angle = 0 - angle;
        }

        // 计算旋转后的新尺寸
        int width = image.getWidth();
        int height = image.getHeight();

        // 计算旋转后的新尺寸（90°和270°需要交换宽高）
        int newWidth = (angle % 180 == 0) ? width : height;
        int newHeight = (angle % 180 == 0) ? height : width;

        // 创建新图像
        BufferedImage rotatedImage = new BufferedImage(
                newWidth,
                newHeight,
                image.getType()
        );

        Graphics2D g2d = rotatedImage.createGraphics();
        AffineTransform transform = new AffineTransform();

        // 移动旋转中心到图像中心
        transform.translate(newWidth / 2.0, newHeight / 2.0);

        // 应用旋转（弧度制）
        double radians = Math.toRadians(angle);
        transform.rotate(radians);

        // 将图像移回原点
        transform.translate(-width / 2.0, -height / 2.0);

        // 绘制并释放资源
        g2d.setTransform(transform);
        g2d.drawImage(image, 0, 0, null);
        g2d.dispose();

        return rotatedImage;
    }

}
